from __future__ import absolute_import, print_function
import os
import os.path as osp
from PIL import Image
from torch.utils.data import Dataset


def read_image(path_img):
    if not osp.exists(path_img):
        raise IOError("The image path does not exist.")
    try:
        img = Image.open(path_img).convert("RGB")
    except:
        raise IOError("Something wrong with reading the image.")
    return img


class PersonImageDataset(Dataset):
    def __init__(self, dataset, transform=None):
        super(PersonImageDataset, self).__init__()
        self.dataset = dataset
        self.transform = transform
        # print('aaaaaaaaaaaaaaaaaaaaaa')
        # print(self.dataset[:5])
        # print('aaaaaaaaaaaaaaaaaaaaaa')

    def __getitem__(self, index):
        path_img, pid, camid = self.dataset[index]
        img = read_image(path_img)

        if self.transform is not None:
            img = self.transform(img)

        return img, path_img, pid, camid

    def __len__(self):
        return len(self.dataset)


if __name__ == "__main__":
    # from data_manager import Market1501
    # market1501 = Market1501(r'D:\reid\DATA4REID')
    # market_trainSet = PersonImageDataset(market1501.train, transform=None)
    # market_querySet = PersonImageDataset(market1501.query, transform=None)
    # market_gallerySet = PersonImageDataset(market1501.gallery, transform=None)
    # print(len(market_trainSet))
    # print(len(market_querySet))
    # print(len(market_gallerySet))
    # print('========== key-vals:')
    # for k, v in market1501.__dict__.items():
    #     print(f"{k}: {v}")

    # from data_manager2 import DukeMTMC
    # dukeMTMC = DukeMTMC(r'D:\reid\DATA4REID')
    # duke_trainSet = PersonImageDataset(dukeMTMC.train, transform=None)
    # duke_querySet = PersonImageDataset(dukeMTMC.query, transform=None)
    # duke_gallerySet = PersonImageDataset(dukeMTMC.gallery, transform=None)
    pass
